{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp pyspark_utils\n",
    "import os\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pyspark Utils\n",
    "\n",
    "Isolate pyspark related code to make pyspark module optional."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from typing import Dict, Tuple, List\n",
    "import os\n",
    "\n",
    "from loguru import logger\n",
    "\n",
    "from pyspark import RDD, SparkContext\n",
    "from pyspark.storagelevel import StorageLevel\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export \n",
    "class Hdfs(object):\n",
    "    \"\"\"Convenience wrapper of hdfs functions for pyspark\"\"\"\n",
    "\n",
    "    def __init__(self, sc=None):\n",
    "        \"\"\" Init Hdfs instance.\n",
    "        Args:\n",
    "          sc: SparkContext instance.\n",
    "        \"\"\"\n",
    "        if sc is None:\n",
    "            sc = SparkContext.getOrCreate()\n",
    "        self.sc = sc\n",
    "        self._hadoop_conf = sc._jsc.hadoopConfiguration()\n",
    "\n",
    "    def list(self, path):\n",
    "        \"\"\"Return files under path\"\"\"\n",
    "        p = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(path)\n",
    "        fs = p.getFileSystem(self._hadoop_conf)\n",
    "        files = fs.listStatus(p)\n",
    "        return [f.getPath().toString() for f in files]\n",
    "\n",
    "    def exists(self, path):\n",
    "        \"\"\"Return path is exists or not\"\"\"\n",
    "        p = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(path)\n",
    "        fs = p.getFileSystem(self._hadoop_conf)\n",
    "        return fs.exists(p)\n",
    "\n",
    "    def remove(self, path):\n",
    "        \"\"\"Remove/delete path\"\"\"\n",
    "        p = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(path)\n",
    "        fs = p.getFileSystem(self._hadoop_conf)\n",
    "        return fs.delete(p, True)\n",
    "\n",
    "    def copy(self, src_path, dst_path):\n",
    "        \"\"\"copy src_path to dst_path\"\"\"\n",
    "        sp = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(src_path)\n",
    "        dp = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(dst_path)\n",
    "        src_fs = sp.getFileSystem(self._hadoop_conf)\n",
    "        dst_fs = dp.getFileSystem(self._hadoop_conf)\n",
    "        return self.sc._gateway.jvm.org.apache.hadoop.fs.FileUtil.copy(\n",
    "            src_fs, sp, dst_fs, dst_path, False, True, self._hadoop_conf)\n",
    "\n",
    "    def copyToLocalFile(self, src_path, dst_path, checksum=True):\n",
    "        \"\"\"copy src_path to dst_path\n",
    "        added by alonso\n",
    "        \"\"\"\n",
    "        sp = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(src_path)\n",
    "        dp = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(dst_path)\n",
    "        fs = sp.getFileSystem(self._hadoop_conf)\n",
    "        fs.setVerifyChecksum(checksum)\n",
    "        fs.copyToLocalFile(sp, dp)\n",
    "\n",
    "    def copyFromLocalFile(self, src_path, dst_path, checksum=True):\n",
    "        \"\"\"added by alonso\"\"\"\n",
    "        sp = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(src_path)\n",
    "        dp = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(dst_path)\n",
    "        fs = dp.getFileSystem(self._hadoop_conf)\n",
    "        fs.setVerifyChecksum(checksum)\n",
    "        fs.copyFromLocalFile(sp, dp)\n",
    "\n",
    "    def mkdir(self, path):\n",
    "        \"\"\"Make directory\"\"\"\n",
    "        p = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(path)\n",
    "        fs = p.getFileSystem(self._hadoop_conf)\n",
    "        return fs.mkdirs(p)\n",
    "\n",
    "    def get_hdfs_rdd(self, path):\n",
    "        return self.sc.textFile(path)\n",
    "\n",
    "    def create(self, path):\n",
    "        p = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(path)\n",
    "        fs = p.getFileSystem(self._hadoop_conf)\n",
    "        return fs.create(p)\n",
    "\n",
    "    def write(self, content, path, data_type=\"str\"):\n",
    "        p = self.sc._gateway.jvm.org.apache.hadoop.fs.Path(path)\n",
    "        fs = p.getFileSystem(self._hadoop_conf)\n",
    "        hdfs_writer = fs.create(p)\n",
    "        if data_type == \"str\":\n",
    "            hdfs_writer.write(bytearray(content, \"utf-8\"))\n",
    "        else:\n",
    "            hdfs_writer.write(content)\n",
    "        hdfs_writer.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "\n",
    "def join_dict_of_rdd(rdd_dict: Dict[str, RDD]) -> RDD:\n",
    "    \"\"\"Join dictionary of RDD, but not a traditional join\n",
    "\n",
    "    1. Stack all RDD together and make as paired RDD\n",
    "    2. reducebyKey, value is list of dict\n",
    "    3. fill all field accorddingly, and return a RDD of dict\n",
    "\n",
    "    Args:\n",
    "        rdd_dict (Dict[str, RDD]): RDD dict\n",
    "\n",
    "    Returns:\n",
    "        RDD: Joined RDD\n",
    "    \"\"\"\n",
    "\n",
    "    # create loss multiplier for inputs\n",
    "    rdd_list = []\n",
    "    loss_multiplier_list = []\n",
    "\n",
    "    def _add_loss_multiplier(inp: dict, problem: str) -> dict:\n",
    "        lm_name = '{}_loss_multiplier'.format(problem)\n",
    "        inp[lm_name] = 1\n",
    "        return inp\n",
    "    for p, rdd in rdd_dict.items():\n",
    "        loss_multiplier_list.append('{}_loss_multiplier'.format(p))\n",
    "        rdd_list.append(\n",
    "            rdd.map(lambda x, p=p: _add_loss_multiplier(x, problem=p))\n",
    "        )\n",
    "\n",
    "    # union rdds\n",
    "    sc: SparkContext = SparkContext.getOrCreate()\n",
    "    all_problem_rdd = sc.union(rdd_list)\n",
    "\n",
    "    # make pair rdd\n",
    "    def _make_pair_rdd(inp_dict: dict) -> Tuple[str, dict]:\n",
    "        if 'record_id' not in inp_dict:\n",
    "            raise KeyError(\n",
    "                \"Chaining problems with & without \"\n",
    "                \"providing 'record_id' in inputs. Received keys: {}\".format(inp_dict.keys()))\n",
    "        return (inp_dict['record_id'], inp_dict)\n",
    "    all_problem_rdd = all_problem_rdd.map(_make_pair_rdd)\n",
    "\n",
    "    # reduce by key, fill out dict correspondingly\n",
    "    loss_multiplier_list_b = sc.broadcast(loss_multiplier_list)\n",
    "\n",
    "    def _merge_dicts(left_dict: dict, right_dict: dict):\n",
    "        left_dict.update(right_dict)\n",
    "        return left_dict\n",
    "\n",
    "    def _add_dummpy_loss_multiplier(inp: dict) -> dict:\n",
    "        # set loss multiplier to inform which problem\n",
    "        # is available in this record\n",
    "        lml = loss_multiplier_list_b.value\n",
    "        for lm in lml:\n",
    "            if lm not in inp:\n",
    "                inp[lm] = 0\n",
    "        return inp\n",
    "    # MEMORY_AND_DISK\n",
    "    all_problem_rdd = all_problem_rdd.persist(\n",
    "        storageLevel=StorageLevel(True, True, False, False))\n",
    "    all_problem_rdd = all_problem_rdd.reduceByKey(_merge_dicts).map(\n",
    "        lambda x: x[1]).map(_add_dummpy_loss_multiplier)\n",
    "\n",
    "    return all_problem_rdd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc : SparkContext= SparkContext.getOrCreate()\n",
    "test_rdd_dict = {\n",
    "    'p1': sc.parallelize([\n",
    "        {'record_id': 0, 'a': 1},\n",
    "        {'record_id': 1, 'a': 2}\n",
    "    ]),\n",
    "    'p2': sc.parallelize([\n",
    "        {'record_id': 0, 'b': 1},\n",
    "        {'record_id': 2, 'b': 2}\n",
    "    ])\n",
    "}\n",
    "rdd = join_dict_of_rdd(test_rdd_dict)\n",
    "assert rdd.collect() == [{'record_id': 0, 'a': 1, 'p1_loss_multiplier': 1, 'b': 1, 'p2_loss_multiplier': 1}, {'record_id': 1, 'a': 2, 'p1_loss_multiplier': 1, 'p2_loss_multiplier': 0}, {'record_id': 2, 'b': 2, 'p2_loss_multiplier': 1, 'p1_loss_multiplier': 0}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_text_file_from_executor(remote_path: str, local_path: str) -> str:\n",
    "    fake_rdd :RDD = SparkContext.getOrCreate().parallelize([1,2,3])\n",
    "    def _read_text_file(x) -> str:\n",
    "        if not os.path.exists(remote_path):\n",
    "            raise FileNotFoundError('{} not found in executor'.format(remote_path))\n",
    "        return [open(remote_path, 'r', encoding='utf8').read()]\n",
    "    file_str_list : List[str]= fake_rdd.mapPartitions(_read_text_file).collect()\n",
    "    if file_str_list:\n",
    "        open(local_path, 'w', encoding='utf8').write(file_str_list[0])\n",
    "    else:\n",
    "        logger.warning('fetch {} to {} failed'.format(remote_path, local_path))\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-06-17 23:46:38.465 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem pyspark_fake_seq_tag, problem type: seq_tag\n",
      "2021-06-17 23:46:38.465 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem pyspark_fake_multi_cls, problem type: multi_cls\n",
      "2021-06-17 23:46:38.466 | INFO     | m3tl.base_params:register_multiple_problems:538 - Adding new problem pyspark_fake_cls, problem type: cls\n",
      "2021-06-17 23:46:38.466 | WARNING  | m3tl.base_params:assign_problem:642 - base_dir and dir_name arguments will be deprecated in the future. Please use model_dir instead.\n",
      "2021-06-17 23:46:38.467 | WARNING  | m3tl.base_params:prepare_dir:361 - bert_config not exists. will load model from huggingface checkpoint.\n",
      "2021-06-17 23:46:44.141 | INFO     | m3tl.utils:set_phase:478 - Setting phase to train\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# hide\n",
    "from m3tl.test_base import PysparkTestBase\n",
    "\n",
    "pyspark_test_base = PysparkTestBase()\n",
    "\n",
    "# create file on executor\n",
    "def create_file_on_executor():\n",
    "    fake_rdd : RDD= SparkContext.getOrCreate().parallelize([1,2,3])\n",
    "    def _make_text_file(x):\n",
    "        open('text.txt', 'w').write('testtest')\n",
    "        return [1]\n",
    "    fake_rdd = fake_rdd.mapPartitions(_make_text_file)\n",
    "    fake_rdd.count()\n",
    "    \n",
    "create_file_on_executor()\n",
    "get_text_file_from_executor('text.txt', 'local_test.txt')\n",
    "assert open('local_test.txt', 'r').read() == 'testtest'\n",
    "os.remove('text.txt')\n",
    "os.remove('local_test.txt')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def repar_rdd(rdd: RDD, rdd_count: int, example_per_par=100000, coalesce_only=True):\n",
    "    \"\"\"\n",
    "    repar rdd based on number of example. if coalesce_only is False and expected\n",
    "    partition is greater than current partition then nothing will happen\n",
    "    \"\"\"\n",
    "    num_partition = rdd.getNumPartitions()\n",
    "    expect_partition = max(1, int(rdd_count / example_per_par))\n",
    "    \n",
    "    if expect_partition < num_partition:\n",
    "        rdd = rdd.coalesce(expect_partition)\n",
    "    elif expect_partition > num_partition and coalesce_only is False:\n",
    "        rdd = rdd.repartition(expect_partition)\n",
    "        \n",
    "    return rdd\n",
    "    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
